#!/usr/bin/env python3

###
# Usage: 
# process single image: prescan_corr_yunda_both_horizonfirst.py img1.fits img1_corrover.fits 1 1 yunda
# process directory: prescan_corr_yunda_both_horizonfirst.py 2024_09_25_14_50_31 . 1 1 yunda

# argv[1]:  input file name or directory name
# argv[2]:  output file name or directory name
# argv[3]:  0:not correct overscan; 1: correct overscan
# argv[4]:  walk=0: just first level directory; walk=1:recursively in satu.
# argv[5]:  CCD type: yunda
###

import sys
import shutil
import numpy as np
from tqdm import trange
from astropy.io import fits
from scipy.signal import savgol_filter
from scipy.stats import trim_mean, sigmaclip
from testlib.imgconfig import load_imgconfig
from glob import glob
import os
import matplotlib
matplotlib.use('Qt5Agg')  # 必须在导入 pyplot 前设置
# from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from PyQt5 import QtWidgets
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
# import matplotlib
# matplotlib.use('Agg')  # 使用非交互模式

filterwidth_v = 250
filterwidth_h = 5

def _smooth_data(data, filterwidth, channel_label=1, direction='Vertical', fig=None, ax=None):
    # 如果未传入 fig/ax，则新建
    if fig is None or ax is None:
        fig, ax = plt.subplots(figsize=(10, 5))
    
    dd = savgol_filter(data, filterwidth * 4 + 1, 3)
    diff = data - dd
    index = np.abs(diff) > sigmaclip(diff, 5, 5)[0].std() * 5
    x = np.arange(len(data))
    dd = np.interp(x, x[~index], data[~index])
    smoothdata = savgol_filter(dd, filterwidth * 2 + 1, 2)
    
    # plt.figure(figsize=(10, 5))
    ax.plot(np.arange(len(data))+1, data, 'b-', label='Data')
    ax.plot(np.arange(len(data))+1, smoothdata, 'r-', label='Smoothed data')
    ax.set_ylim(data.mean() - 4 * data.std(), data.mean() + 4 * data.std())
    # plt.title(f'Overscan Smoothing. Window length={filterwidth*2+1} (Channel {channel_label})')
    ax.set_title(f'Channel {channel_label}', fontsize=10, pad=5)
    ax.tick_params(axis='both', labelsize=8)

    ax.legend()
    ax.grid(True)
    # 如果传入的是外部 fig/ax，则不自动显示
    if fig is not None and ax is not None:
        fig.canvas.draw()  # 仅渲染，不阻塞

    return smoothdata

def extract_overscan_yunda(image, os_x1, os_x2, y1, y2, filterwidth=5, axis_label=1, channel_label=1, direction='Vertical', fig=None, ax=None):
    if os_x1 >= os_x2:
        raise Exception('os_x1 must be smaller than os_x2')
    if y1 >= y2:
        raise Exception('y1 must be smaller than y2')
    data = trim_mean(image[y1:y2, os_x1:os_x2], 0.1, axis=axis_label)
    # np.savetxt(f'overscan_channel{channel_label}_{direction}.txt', data.reshape((data.shape[0],1)), fmt='%.4f')
    data = _smooth_data(data, filterwidth, channel_label=channel_label, direction=direction, fig=fig, ax=ax)
    # print('overscan shape:',data.shape)
    return data

def subtract_overscan(image, x0, nx, y0, ny, os_x1, os_x2):
    v = extract_overscan_yunda(image, os_x1, os_x2, y0, y0 + ny, filterwidth=3)
    # v = np.repeat(v, nx).reshape((ny, nx))
    image[y0:y0 + ny, x0:x0 + nx] -= v.reshape((v.shape[0], 1))
    return image

def crop_overscan(image, imgconf):
    output = np.zeros(imgconf.mos_shape, dtype=image.dtype)
    for chan in range(imgconf.nchan):
        x1 = imgconf.raw_x0[chan]
        y1 = imgconf.raw_y0[chan]
        cimage = image[y1:y1 + imgconf.raw_ny[chan], x1:x1 + imgconf.raw_nx[chan]]

        if imgconf.raw_xflip[chan]:
            x1 = imgconf.raw_nx[chan] - imgconf.x1[chan] - imgconf.nx[chan]
        else:
            x1 = imgconf.x1[chan]
        if imgconf.raw_yflip[chan]:
            y1 = imgconf.raw_ny[chan] - imgconf.y1[chan] - imgconf.ny[chan]
        else:
            y1 = imgconf.y1[chan]
        cimage = cimage[y1:y1 + imgconf.ny[chan], x1:x1 + imgconf.nx[chan]]

        if imgconf.mos_xflip[chan]:
            cimage = cimage[:, ::-1]
        if imgconf.mos_yflip[chan]:
            cimage = cimage[::-1, :]
        x1 = imgconf.mos_x0[chan]
        y1 = imgconf.mos_y0[chan]
        output[y1:y1 + imgconf.ny[chan], x1:x1 + imgconf.nx[chan]] = cimage
    return output

def corr_prescan_both(image, imgconf):
    imgconf.prescan_x1 = [0] * 16   # 通道内prescan x向起始位置
    imgconf.prescan_x2 = [26] * 16   # 通道内prescan x向终止位置
    
    # 创建Qt应用（如果不存在）
    app = QtWidgets.QApplication.instance()
    if app is None:
        app = QtWidgets.QApplication([])
    
    fig_horizon = plt.figure(figsize=(24, 15))
    # gs_horizon = fig_horizon.add_gridspec(4, 4)  # 4行4列的网格布局
    gs_horizon = gridspec.GridSpec(4, 4, hspace=0.3, wspace=0.15)
    # fig_horizon.show()
    
    fig_vertical = plt.figure(figsize=(24, 15))
    # gs_vertical = fig_vertical.add_gridspec(4, 4)  # 4行4列的网格布局
    gs_vertical = gridspec.GridSpec(4, 4, hspace=0.3, wspace=0.15)
    
    
    for chan in range(imgconf.nchan):
        
        x0 = imgconf.raw_x0[chan]
        y0 = imgconf.raw_y0[chan]
        if imgconf.raw_xflip[chan]:
            x1 = x0 + imgconf.raw_nx[chan] - imgconf.x1[chan] - imgconf.nx[chan]
            os1 = x0 + imgconf.raw_nx[chan] - imgconf.prescan_x2[chan]-1 +5
            os2 = x0 + imgconf.raw_nx[chan] - imgconf.prescan_x1[chan]-1 -5
        else:
            x1 = x0 + imgconf.x1[chan]
            os1 = x0 + imgconf.prescan_x1[chan] +5
            os2 = x0 + imgconf.prescan_x2[chan] -5
        if imgconf.raw_yflip[chan]:
            y1 = y0 + imgconf.raw_ny[chan] - imgconf.y1[chan] - imgconf.ny[chan]
        else:
            y1 = y0 + imgconf.y1[chan]
        
        # correct hirizontal overscan
        ax_horizon = fig_horizon.add_subplot(gs_horizon[chan // 4, chan % 4])  # 每行4个子图，共4行
        col_os = extract_overscan_yunda(image, x0, x0 + imgconf.raw_nx[chan], y0, y0 + imgconf.raw_ny[chan], filterwidth=filterwidth_h, axis_label=0, channel_label=str(chan+1), direction='Horizontal', fig=fig_horizon, ax=ax_horizon)
        image[y0:y0+imgconf.raw_ny[chan], x0:x0 + imgconf.raw_nx[chan]] -= col_os
        
        # correct vertical prescan
        ax_vertical = fig_vertical.add_subplot(gs_vertical[chan // 4, chan % 4])
        row_os = extract_overscan_yunda(image, os1, os2, y0, y0 + imgconf.raw_ny[chan], filterwidth=filterwidth_v, axis_label=1, channel_label=str(chan+1), direction='Vertical', fig=fig_vertical, ax=ax_vertical)
        # v = np.repeat(v, nx).reshape((ny, nx))
        image[y0:y0+imgconf.raw_ny[chan], x0:x0 + imgconf.raw_nx[chan]] -= row_os.reshape((row_os.shape[0], 1))
        
    # plt.title()
    # plt.show()
    
    # 显示图形
    fig_horizon.suptitle(f'Horizontal Overscan Smoothing (Sav_Gol width={int(filterwidth_h*2+1)})', y=0.95, fontsize=12)
    fig_vertical.suptitle(f'Vertical Prescan Smoothing (Sav_Gol width={int(filterwidth_v*2+1)})', y=0.95, fontsize=12)
    # fig_horizon.tight_layout(rect=[0, 0, 1, 0.95])
    # fig_vertical.tight_layout(rect=[0, 0, 1, 0.95])
    

    fig_horizon.savefig(f"horizontal_overscan_smoothing_savgolwidth{int(filterwidth_h*2+1)}.png", dpi=300, bbox_inches='tight')
    fig_vertical.savefig(f"vertical_prescan_smoothing_savgolwidth{int(filterwidth_v*2+1)}.png", dpi=300, bbox_inches='tight')
    
    # fig_horizon.show()
    # fig_vertical.show()
    # # 启动事件循环
    # app.exec_()
    plt.close(fig_horizon)
    plt.close(fig_vertical)
    
    return image

def do_overscan(flist, outlist, imgconf, crop=True, croponly=False):
    print(imgconf)
    if croponly:
        crop = True
    # file name
    n = len(flist)
    if len(outlist) != n:
        raise Exception('output name list is not as long as input list')
    # do overscan correction
    for i in trange(n, desc='image cropping/overscan correction'):
        if imgconf.overscan == 1:
            with fits.open(flist[i]) as img:
                data = img[0].data.astype(np.float32)
                if croponly is False and imgconf.type != 'cmos':
                    data = corr_prescan_both(data, imgconf)
                if crop:
                    data = crop_overscan(data, imgconf)
                hdu = fits.PrimaryHDU(data, header=img[0].header)
                hdu.header['ORIGNAME'] = flist[i]
                hdu.writeto(outlist[i], overwrite=True)
    return


if __name__ == '__main__':
    input = sys.argv[1]        # input file name or directory name
    outd = sys.argv[2]         # output file name or directory name
    corr = int(sys.argv[3])    # 0:not correct overscan; 1: correct overscan
    walk = int(sys.argv[4])    # walk=0: just first level directory; walk=1:recursively in satu.
    ctype = sys.argv[5]        # CCD type: xiguang,xiguang_TTD,44ccdmsc
    imgconf = load_imgconfig(ctype)
    if corr==1:
        crop,croponly = True,False
    else:
        crop,croponly = True,True
    if input.endswith(".fits"):
        flist = [input]
        outlist = [outd]
        do_overscan(flist, outlist, imgconf, crop, croponly)
    elif walk==1:
        flist = []
        outlist = []
        for r,d,fs in os.walk(input):
            for f in fs:
                if f.endswith(".fits") and (not f.endswith("_prescan_corr.fits")):
                    flist.append(r+"/"+f)
                    outlist.append(r+"/"+f.replace(".fits","_both_prescan_corr.fits"))
        do_overscan(flist, outlist, imgconf, crop, croponly)  #imgconfig:xiguang,xiguang_TTD  crop 0,1  croponly 0(crop&subtract overscan),1(not subtract overscan,only crop)
    else:
        fs = glob(input+"/*.fits")
        flist = []
        outlist = []
        for f in fs:
            fname = os.path.basename(f)
            if f.endswith(".fits") and (not f.endswith("_corr.fits")):
                flist.append(f)
                # outlist.append(outd+"/"+fname.replace(".fits","._both_corr.fits"))
                outlist.append(os.path.join(outd, fname.replace(".fits","._both_prescan_corr.fits")))
        do_overscan(flist, outlist, imgconf, crop, croponly)